package com.supreme.secretserver.component;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.supreme.common.constant.ConsoleColors;
import com.supreme.commonboot.utils.AjaxResp;
import com.supreme.secretserver.annotation.Encrypt;
import com.supreme.secretserver.props.EncryptProps;
import com.supreme.secretserver.secret.SecretParser;
import com.supreme.secretserver.utils.AESUtils;
import com.supreme.secretserver.utils.RSAUtils;
import org.apache.commons.lang3.StringUtils;
import org.springframework.core.MethodParameter;
import org.springframework.http.MediaType;
import org.springframework.http.converter.HttpMessageConverter;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.web.bind.annotation.ControllerAdvice;
import org.springframework.web.servlet.mvc.method.annotation.ResponseBodyAdvice;

import javax.annotation.Resource;
import java.util.Objects;
import java.util.Optional;

@ControllerAdvice
public class EncryptResponse implements ResponseBodyAdvice<AjaxResp> {

    private final ObjectMapper objectMapper = new ObjectMapper();

    @Resource
    private EncryptProps encryptProps;

    @Override
    public boolean supports(MethodParameter methodParameter, Class<? extends HttpMessageConverter<?>> clazz) {
        if (methodParameter.hasMethodAnnotation(Encrypt.class)) {
            Encrypt encryptAnno = methodParameter.getMethodAnnotation(Encrypt.class);
            Boolean encryptEnable = Optional.ofNullable(encryptAnno)
                    .filter(Encrypt::enable)
                    .filter(encrypt -> StringUtils.equalsAny(encrypt.type(), SecretParser.AES, SecretParser.RSA))
                    .isPresent();
            ConsoleColors.println(ConsoleColors.CYAN_BOLD_BRIGHT, "EncryptResponse encryptEnable: " + encryptEnable);
            return encryptEnable;
        }

        return false;
    }

    @Override
    public AjaxResp beforeBodyWrite(AjaxResp ajaxResp, MethodParameter methodParameter, MediaType mediaType, Class<? extends HttpMessageConverter<?>> clazz, ServerHttpRequest serverHttpRequest, ServerHttpResponse serverHttpResponse) {
        if (Objects.nonNull(ajaxResp)) {
            Encrypt methodAnnotation = methodParameter.getMethodAnnotation(Encrypt.class);
            Optional.ofNullable(methodAnnotation)
                    .filter(Encrypt::enable)
                    .filter(encrypt -> StringUtils.isNotBlank(encrypt.type()))
                    .ifPresent(encrypt -> {
                        try {
                            String secretType = encrypt.type();
                            Object dataObject = ajaxResp.getData();
                            byte[] dataBytes;
                            if (dataObject instanceof String) {
                                dataBytes = dataObject.toString().getBytes();
                            } else {
                                dataBytes = objectMapper.writeValueAsBytes(dataObject);
                            }

                            if (SecretParser.AES.equals(secretType)) {
                                String secretKey = encrypt.secretKey();
                                byte[] encryptKeyBytes = Optional.ofNullable(secretKey)
                                        .filter(StringUtils::isNotBlank)
                                        .map(String::getBytes)
                                        .orElse(encryptProps.getAesKeyBytes());
                                String encryptBody = AESUtils.encrypt(dataBytes, encryptKeyBytes);
                                ajaxResp.setData(encryptBody);
                            } else if (SecretParser.RSA.equals(secretType)) {
                                String publicKey = encrypt.publicKey();
                                byte[] encryptKeyBytes = Optional.ofNullable(publicKey)
                                        .filter(StringUtils::isNotBlank)
                                        .map(String::getBytes)
                                        .orElse(encryptProps.getRsaPubKey().getBytes());
                                String encryptBody = RSAUtils.encryptByPublicKey(encryptKeyBytes, dataBytes);
                                ajaxResp.setData(encryptBody);
                            }
                        } catch (Exception e) {
                            e.printStackTrace();
                        }
                    });
        }

        return ajaxResp;
    }
}
